
from astropy import units as u
from netorchestr.envir.base import OMessage, OLink, OGate, OModule, ONet

class Txc(OModule):
    def __init__(self, name:str):
        super().__init__(name)

    def initialize(self):
        if self.name == 'tic':
            msg = OMessage(0, 'tic', 'toc', 'tictocMsg', '0')
            out_gate = [g for g in self.gates if g.name == 'out']
            self.send_msg(msg, out_gate[0])

    def recv_msg(self, msg:OMessage, gate:OGate):
        msg.id = f"{int(msg.id)+1}"
        out_gate = [g for g in self.gates if g.name == 'out']
        self.send_msg(msg, out_gate[0])

if __name__ == '__main__':
    net = ONet('Tictoc')
    tic = Txc('tic')
    toc = Txc('toc')
    
    gate_tic = OGate('out', tic)
    gate_toc = OGate('out', toc)
    
    tic.gates[gate_tic] = (OLink('tic-toc', 20.0*u.ms), gate_toc)
    toc.gates[gate_toc] = (OLink('tic-toc', 100.0*u.ms), gate_tic)
    
    net.add_module(tic)
    net.add_module(toc)
    
    net.run(1000*u.ms)
